#include <sds/sds.h>

#include <cluster.h>
#include <matrix.h>

#include <random.h>

#include <vector>
#include <iostream>
#include <iomanip>

using namespace sds;
using namespace std;

Random rng;

int main(int argc, const char* argv[]) {
  SFile clufile(argv[2]);
  SDataView<unsigned char> clu =
    clufile.getDataset(0).getDataView<unsigned char>();

  SFile reffile(argv[3]);
  SDataView<unsigned char> ref =
    reffile.getDataset(0).getDataView<unsigned char>();

  vector<SInteger> size = ref.getSize();

  SFile srcfile(argv[1]);
  SDataView<Point> src(size);

  for (int k = 0; k < CH; ++k) {
    SDataView<unsigned char> view =
      srcfile.getDataset(k).getDataView<unsigned char>();
    for (int i = 0; i < size[0]; ++i) {
      for (int j = 0; j < size[1]; ++j) {
	((Point&)src[i][j])[k] = (double)view[i][j];
      }    
    }
  }

  vector<Cluster> clusters(256);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (clu[i][j] != 0) {
	clusters[clu[i][j]] += src[i][j];
      }
    }
  }

  vector<int> index;
  vector<Vector> expected;
  vector<Matrix> inverse;
  vector<double> lndet;

  for (int k = 1; k < 256; ++k) {
    if (clusters[k].size() == 0) continue;
    index.push_back(k);
    Vector exp(CH);
    for (int i = 0; i < CH; ++i) {
      exp(i) = clusters[k].expected(i);
    }
    expected.push_back(exp);
    Matrix cov(CH, CH);
    for (int i = 0; i < CH; ++i) {
      for (int j = 0; j < CH; ++j) {
	cov(i, j) = clusters[k].covariance(i, j);
      }
      cov(i, i) += 0.00001;
    }
    lndet.push_back(log(fabs(cov.determinant())));
    inverse.push_back(cov.inverse());
  }

  SDataView<unsigned char> trg(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Vector current(CH);
      for (int l = 0; l < CH; ++l) {
	current(l) = ((Point&)src[i][j])[l];
      }
      int minPos = -1; 
      double minVal = 0.0;
      for (int k = 0; k < index.size(); ++k) {
	Vector diff = current - expected[k];
	double dist = diff * inverse[k] * diff + lndet[k];
	if (minPos == -1 || minVal > dist) {
	  minPos = k;
	  minVal = dist;
	}
      }
      trg[i][j] = minPos;
    }
  }

  int cr[256][256], r[256], s[256], fs[256];
  
  for (int i = 0; i < index.size(); ++i) {
    for (int j = 0; j < 256; ++j) {
      cr[i][j] = 0;
    }
    r[i] = 0; s[i] = 0; fs[i] = 0;
  }
  
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      ++fs[trg[i][j]];
      if (ref[i][j] == 0) continue; 
      ++cr[trg[i][j]][ref[i][j]];
      ++r[ref[i][j]];
      ++s[trg[i][j]];
    }
  }

  vector<unsigned char> d[256];
  vector<Cluster> divclusters[256];

  unsigned char c[256];
  for (int i = 0; i < index.size(); ++i) {
    c[i] = 0;
    if (fs[i] >= 1000 * s[i]) continue;
    std::cout << "Relevant colors: ";
    for (int j = 0; j < 256; ++j) {
      if (cr[i][j] >= 0.75 * s[i]) {
	std::cout << '*';
 	c[i] = j;
      }
      if (cr[i][j] >= 0.25 * s[i]) {
	std::cout << j << ' ';
 	d[i].push_back(j);
	divclusters[i].push_back(Cluster());
      }
    }
    std::cout << std::endl;
  }

  std::cout << index.size() << std::endl;
  int u = 0;
  std::vector<int> t(5, 0);
  for (int i = 0; i < index.size(); ++i) {
    if (c[i] == 0) {
      ++t[d[i].size()];
    } else {
      ++u;
    }
  }

  std::cout << argv[2] << std::endl;
  std::cout << "U" << ' ' << u << std::endl;
  for (int i = 0; i < t.size(); ++i) {
    std::cout << i << ' ' << t[i] << std::endl;
  }


  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (c[trg[i][j]] == 0) {
	for (int k = 0; k < d[trg[i][j]].size(); ++k) {
          if (ref[i][j] == d[trg[i][j]][k]) {
            divclusters[trg[i][j]][k] += src[i][j];
          }
	}
      }
    }
  }

  vector<Vector> divexpected[256];
  vector<Matrix> divinverse[256];
  vector<double> divlndet[256];

  for (int l = 0; l < index.size(); ++l) {
    for (int k = 0; k < d[l].size(); ++k) {
      Vector exp(CH);
      for (int i = 0; i < CH; ++i) {
	exp(i) = divclusters[l][k].expected(i);
      }
      divexpected[l].push_back(exp);
      Matrix cov(CH, CH);
      for (int i = 0; i < CH; ++i) {
	for (int j = 0; j < CH; ++j) {
	  cov(i, j) = divclusters[l][k].covariance(i, j);
	}
	cov(i, i) += 0.0001;
      }
      divlndet[l].push_back(fabs(log(cov.determinant())));
      divinverse[l].push_back(cov.inverse());
    }
  }
  
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (c[trg[i][j]] != 0) {
	trg[i][j] = c[trg[i][j]];
      } else {
	Vector current(CH);
	for (int l = 0; l < CH; ++l) {
	  current(l) = ((Point&)src[i][j])[l];
	}
	int minPos = -1; 
	double minVal = 0.0;
	for (int k = 0; k < d[trg[i][j]].size(); ++k) {
	  Vector diff = current - divexpected[trg[i][j]][k];
	  double dist = divlndet[trg[i][j]][k] + 
            diff * divinverse[trg[i][j]][k] * diff;
	  if (minPos == -1 || minVal > dist) {
	    minPos = d[trg[i][j]][k];
	    minVal = dist;
	  }
	}
	trg[i][j] = minPos == -1 ? 0 : minPos;
      }
    }
  }

  SFile trgfile(argv[4], SFile::CREATE);
  trgfile.createDataset("clasified", UINT8, size).
    setDataView<unsigned char>(trg);
  return 0;
}
